import torch
from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
from PIL import Image
import numpy as np
import requests
import base64
import json

# 1. 增强的文生图生成器，支持多种模型
class TextToImageGenerator:
    def __init__(self, model_type="stable_diffusion", model_name="stabilityai/stable-diffusion-2-1"):
        self.model_type = model_type
        
        if model_type == "stable_diffusion":
            self._init_stable_diffusion(model_name)
        elif model_type == "seedream":
            self._init_seedream()
    
    def _init_stable_diffusion(self, model_name):
        """初始化 Stable Diffusion 模型"""
        self.pipe = StableDiffusionPipeline.from_pretrained(
            model_name,
            torch_dtype=torch.float16,
            safety_checker=None
        )
        self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
        self.pipe = self.pipe.to("cuda" if torch.cuda.is_available() else "cpu")
        self.pipe.enable_attention_slicing()
    
    def _init_seedream(self):
        """初始化 Seedream 3.0 API 配置"""
        # 火山引擎 API 配置
        self.seedream_config = {
            "api_url": "https://visual.volcengineapi.com",
            "access_key": "YOUR_ACCESS_KEY",  # 需要替换为实际的 API Key
            "secret_key": "YOUR_SECRET_KEY",   # 需要替换为实际的 Secret Key
            "service": "cv",
            "version": "2022-08-31",
            "action": "SeedreamGenerate"
        }
    
    def generate_image(self, prompt, negative_prompt="", steps=28, guidance_scale=7.5, width=512, height=512):
        """根据模型类型生成图像"""
        if self.model_type == "stable_diffusion":
            return self._generate_with_sd(prompt, negative_prompt, steps, guidance_scale, width, height)
        elif self.model_type == "seedream":
            return self._generate_with_seedream(prompt, negative_prompt, width, height)
    
    def _generate_with_sd(self, prompt, negative_prompt, steps, guidance_scale, width, height):
        """使用 Stable Diffusion 生成图像"""
        image = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=steps,
            guidance_scale=guidance_scale,
            width=width,
            height=height
        ).images[0]
        return image
    
    def _generate_with_seedream(self, prompt, negative_prompt, width=1024, height=1024):
        """使用 Seedream 3.0 生成图像"""
        # Seedream API 请求参数
        request_data = {
            "req_key": "seedream_generate",
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "width": width,
            "height": height,
            "model_version": "seedream_v3.0",
            "num_images": 1,
            "quality": "high",
            "style": "realistic"  # 可选：realistic, anime, artistic
        }
        
        try:
            # 发送 API 请求
            response = self._call_seedream_api(request_data)
            
            if response["status"] == "success":
                # 解码 base64 图像
                image_data = base64.b64decode(response["data"]["images"][0])
                image = Image.open(io.BytesIO(image_data))
                return image
            else:
                print(f"Seedream API 错误: {response['message']}")
                return None
                
        except Exception as e:
            print(f"Seedream 生成失败: {str(e)}")
            return None
    
    def _call_seedream_api(self, request_data):
        """调用 Seedream API"""
        import hashlib
        import hmac
        import time
        from urllib.parse import urlencode
        
        # 构建请求签名
        timestamp = str(int(time.time()))
        headers = {
            "Content-Type": "application/json",
            "X-Date": timestamp,
            "Authorization": self._generate_auth_header(request_data, timestamp)
        }
        
        # 发送请求
        response = requests.post(
            f"{self.seedream_config['api_url']}/seedream/generate",
            headers=headers,
            json=request_data,
            timeout=30
        )
        
        return response.json()
    
    def _generate_auth_header(self, request_data, timestamp):
        """生成火山引擎 API 认证头"""
        # 这里需要实现火山引擎的签名算法
        # 具体实现请参考火山引擎 API 文档
        pass

# 2. 增强的提示词优化器，支持 Seedream 特性
def optimize_prompt_for_model(transcribed_text, model_type="stable_diffusion"):
    """根据模型类型优化提示词"""
    import jieba.analyse
    
    # 提取关键词
    keywords = jieba.analyse.extract_tags(transcribed_text, topK=5, withWeight=False)
    
    if model_type == "stable_diffusion":
        return _optimize_for_sd(keywords)
    elif model_type == "seedream":
        return _optimize_for_seedream(keywords, transcribed_text)

def _optimize_for_sd(keywords):
    """为 Stable Diffusion 优化提示词"""
    quality_tags = "masterpiece, best quality, 4k, detailed"
    prompt = f"{quality_tags}, {', '.join(keywords)}"
    negative_prompt = "worst quality, low quality, deformed, mutation, blurry"
    return prompt, negative_prompt

def _optimize_for_seedream(keywords, original_text):
    """为 Seedream 3.0 优化提示词"""
    # Seedream 3.0 支持中文提示词，保持原文
    base_prompt = original_text
    
    # 添加 Seedream 专用质量标签
    quality_enhancer = "高质量, 精细细节, 专业摄影, 4K分辨率"
    prompt = f"{base_prompt}, {quality_enhancer}"
    
    # Seedream 负面提示词
    negative_prompt = "低质量, 模糊, 变形, 噪点, 水印"
    
    return prompt, negative_prompt

# 3. 模型选择器
def create_generator(model_choice="auto"):
    """智能选择最适合的模型"""
    if model_choice == "auto":
        # 自动选择：优先使用 Seedream（如果配置了 API），否则使用 SD
        try:
            generator = TextToImageGenerator("seedream")
            print("使用 Seedream 3.0 模型")
            return generator, "seedream"
        except:
            print("Seedream 不可用，回退到 Stable Diffusion")
            generator = TextToImageGenerator("stable_diffusion")
            return generator, "stable_diffusion"
    
    elif model_choice == "seedream":
        return TextToImageGenerator("seedream"), "seedream"
    elif model_choice == "sd":
        return TextToImageGenerator("stable_diffusion"), "stable_diffusion"

# 4. 主流程整合
if __name__ == "__main__":
    print("选择文生图模型:")
    print("1. Seedream 3.0 (火山引擎)")
    print("2. Stable Diffusion (本地)")
    print("3. 自动选择")
    
    choice = input("请输入选项 (1-3): ").strip()
    
    model_map = {"1": "seedream", "2": "sd", "3": "auto"}
    model_choice = model_map.get(choice, "auto")
    
    # 创建生成器
    generator, selected_model = create_generator(model_choice)
    
    # 模拟语音识别结果（实际应该从语音识别模块获取）
    text = "一只可爱的小猫坐在花园里，阳光明媚，超现实主义风格"
    
    # 优化提示词
    prompt, negative_prompt = optimize_prompt_for_model(text, selected_model)
    print(f"优化后的提示词: {prompt}")
    print(f"负面提示词: {negative_prompt}")
    
    # 生成图像
    print("正在生成图像...")
    
    if selected_model == "seedream":
        # Seedream 3.0 推荐参数
        image = generator.generate_image(
            prompt,
            negative_prompt=negative_prompt,
            width=1024,
            height=1024
        )
    else:
        # Stable Diffusion 参数
        image = generator.generate_image(
            prompt,
            negative_prompt=negative_prompt,
            steps=28,
            guidance_scale=7.5,
            width=512,
            height=512
        )
    
    if image:
        # 保存结果
        filename = f"generated_image_{selected_model}.png"
        image.save(filename)
        print(f"图像已保存为 {filename}")
        
        # 显示图像信息
        print(f"图像尺寸: {image.size}")
        print(f"使用模型: {selected_model}")
    else:
        print("图像生成失败")

# 5. Seedream 3.0 特色功能
class SeedreamEnhancer:
    """Seedream 3.0 特色功能增强"""
    
    @staticmethod
    def style_transfer(base_prompt, style="国潮"):
        """风格迁移"""
        style_prompts = {
            "国潮": "中国传统文化元素, 现代设计, 故宫红, 祥云纹样",
            "赛博朋克": "霓虹灯, 未来科技, 机械美学, 蓝紫色调",
            "水墨画": "中国水墨画风格, 写意, 留白, 墨色渲染",
            "油画": "欧洲古典油画, 厚重笔触, 明暗对比, 古典美"
        }
        
        style_tag = style_prompts.get(style, "")
        return f"{base_prompt}, {style_tag}"
    
    @staticmethod
    def lighting_enhance(prompt, lighting="自然光"):
        """光照增强"""
        lighting_tags = {
            "自然光": "自然光照, 柔和阴影",
            "戏剧光": "戏剧性光照, 强烈对比",
            "逆光": "逆光拍摄, 轮廓光",
            "黄金时刻": "黄金时刻, 暖色调光线"
        }
        
        light_tag = lighting_tags.get(lighting, "")
        return f"{prompt}, {light_tag}"